package com.gxc.tree;

import java.util.HashMap;
import java.util.Map;
import java.util.Stack;

public class KthSmallest {

    public int kthSmallest(TreeNodeK root, int k) {
        Stack<TreeNodeK> stack = new Stack<>();
        TreeNodeK pop = null;
        while (!stack.isEmpty() || root!=null){
            while (root != null) {
                stack.push(root);
                root = root.left;
            }
            pop = stack.pop();
            k--;
            if (k == 0) {
                break;
            }
            root = pop.right;
        }
        return pop.val;
    }

    public static void main(String[] args) {
        TreeNodeK root = new TreeNodeK(3);
        TreeNodeK a1 = new TreeNodeK(1);
        TreeNodeK a2 = new TreeNodeK(4);
        root.left = a1;
        root.right = a2;
        TreeNodeK a3 = new TreeNodeK(5);
        a1.right = a3;
        System.out.println(kthSmallest2(root, 1));
    }

    public static int kthSmallest2(TreeNodeK root, int k) {
        MyBst myBst = new MyBst(root);
        return myBst.kthSmallest(k);
    }
}

class TreeNodeK {
    int val;
    TreeNodeK left;
    TreeNodeK right;
    TreeNodeK() {}
    TreeNodeK(int val) { this.val = val; }
    TreeNodeK(int val, TreeNodeK left, TreeNodeK right) {
        this.val = val;
        this.left = left;
        this.right = right;
    }
}

class MyBst {
    TreeNodeK root;
    Map<TreeNodeK, Integer> nodeNum;

    public MyBst(TreeNodeK root) {
        this.root = root;
        nodeNum = new HashMap<>();
        countNodeNum(root);
    }

    private int countNodeNum(TreeNodeK node) {
        if (node == null) return 0;
        int left = countNodeNum(node.left);
        int right = countNodeNum(node.right);
        int nodeSum = left + right + 1;
        nodeNum.put(node, nodeSum);
        return nodeSum;
    }

    public int kthSmallest(int k) {
        TreeNodeK node = this.root;
        while (node != null) {
            Integer left = getNodeNum(node.left);
            if (left < k -1 ) {
                node = node.right;
                k = k - (left +1);
            } else if (left == k-1) {
                break;
            } else {
                node = node.left;
            }
        }
        return node.val;
    }

    // 获取以node为根结点的子树的结点数
    private int getNodeNum(TreeNodeK node) {
        return nodeNum.getOrDefault(node, 0);
    }
}
